import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import transforms
import os
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import matplotlib.pyplot as plt
from model import CNN

base_dir = '新预处理'
image_train_dir = os.path.join(base_dir, 'train')
image_val_dir = os.path.join(base_dir, 'val')
image_test_dir = os.path.join(base_dir, 'test')
excel_train = os.path.join(base_dir, 'train.xlsx')
excel_val = os.path.join(base_dir, 'val.xlsx')
excel_test = os.path.join(base_dir, 'test.xlsx')

class CustomDataset(Dataset):
    def __init__(self, annotations_file, image_dir, transform=None):
        self.df = pd.read_excel(annotations_file)
        self.image_dir = image_dir
        self.transform = transform

    def __len__(self):
        return len(self.df)-1

    def __getitem__(self, idx):
        filename = f"{idx + 1}.jpg"
        img_path = os.path.join(self.image_dir, filename)
        image = Image.open(img_path).convert('RGB')
        label_number = self.df.iloc[idx, 13]
        label = label_number
        label = torch.tensor(label, dtype=torch.uint8) #dtype=torch.float32 torch.uint8
        if self.transform:
            image = self.transform(image)
        return image, label

#数据增强
train_transform = transforms.Compose([transforms.Resize((256, 256)),  # 统一缩放
                                transforms.RandomRotation(90),     #随机旋转
                                transforms.RandomHorizontalFlip(0.5),
                                transforms.RandomVerticalFlip(0.5),
                                transforms.ToTensor(),        # 转换为tensor
                                transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                                     std=[0.5, 0.5, 0.5])])
val_transform = transforms.Compose([transforms.Resize((256, 256)),  # 统一缩放
                                transforms.ToTensor(),        # 转换为tensor
                                transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                                     std=[0.5, 0.5, 0.5])])

train_ds = CustomDataset(annotations_file=excel_train, image_dir=image_train_dir, transform=train_transform)
val_ds = CustomDataset(annotations_file=excel_val, image_dir=image_val_dir, transform=val_transform)

# dataloader
batch_size = 16  # batch 批次
train_loader = torch.utils.data.DataLoader(train_ds, batch_size=batch_size, shuffle=True, drop_last=True)
val_loader = torch.utils.data.DataLoader(val_ds, batch_size=batch_size)

#可视化某一批数据
'''train_img,train_label=next(iter(train_loader))   #iter迭代器，可以用来便利trainloader里面每一个数据，这里只迭代一次来进行可视化
fig, axes = plt.subplots(2, 4, figsize=(10, 10))
axes_list = []
#输入到网络的图像
for i in range(axes.shape[0]):
    for j in range(axes.shape[1]):
        axes[i, j].imshow(train_img[i*4+j,0,:,:],cmap="gray")    #这里画出来的就是我们想输入到网络里训练的图像，与之对应的标签用来进行最后分类结果损失函数的计算
        axes[i, j].axis("off")
#对应的标签
print(train_label)
plt.show()'''

# 加载预训练模型resnet50
model = torchvision.models.resnet50(pretrained=True)
for param in model.parameters():
    param.requires_grad = False
# 修改结果特征
model.fc = nn.Linear(model.fc.in_features, 5)
for i in range(3,5):
    for param in getattr(model, f"layer{i}").parameters():
        param.requires_grad = True

# 加载预训练模型vgg16
'''model = torchvision.models.vgg16_bn(pretrained=True)
model.classifier[6] = nn.Linear(model.classifier[6].in_features, 5)
for i, param in enumerate(model.parameters()):
    print(i)
    if i < 33:      # 前面一些参数冻结
        param.requires_grad = False
print(model)'''

#model = torchvision.models.inception_v3(pretrained=True)
#vit swin

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)

'''device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = nn.DataParallel(CNN(in_ch=3)).to(device)'''

# 继续训练
optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()

# 定义训练过程
def train(epoch, model, train_loader, test_loader):
    if epoch == 5:
        print("now!")
    correct = 0
    total = 0
    running_loss = 0

    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        with torch.no_grad():
            y_pred = torch.argmax(y_pred, dim=1)
            correct += (y_pred == y).sum().item()
            total += y.size(0)
            running_loss += loss.item()

    epoch_loss = running_loss / len(train_loader.dataset)
    epoch_acc = correct / total

    # 测试过程
    test_correct = 0
    test_total = 0
    test_running_loss = 0
    with torch.no_grad():
        for x, y in test_loader:
            x, y = x.to(device), y.to(device)
            y_pred = model(x)
            loss = loss_fn(y_pred, y)
            y_pred = torch.argmax(y_pred, dim=1)
            test_correct += (y_pred == y).sum().item()
            test_total += y.size(0)
            test_running_loss += loss.item()
    test_epoch_loss = test_running_loss / len(test_loader.dataset)
    test_epoch_acc = test_correct / test_total

    print('epoch', epoch,
          'loss', round(epoch_loss, 3),
          'accuracy', round(epoch_acc, 3),
          'test_loss', round(test_epoch_loss, 3),
          'test_accuracy', round(test_epoch_acc, 3))
    return epoch_loss, epoch_acc, test_epoch_loss, test_epoch_acc

epoch = 4
all_epoch_loss = []
all_epoch_acc = []
all_val_loss = []
all_val_acc = []
for i in range(epoch):
    epoch_loss, epoch_acc, test_epoch_loss, test_epoch_acc = train(i, model, train_loader, val_loader)
    all_epoch_loss.append(epoch_loss)
    all_val_loss.append(test_epoch_loss)
    all_epoch_acc.append(epoch_acc)
    all_val_acc.append(test_epoch_acc)

x = range(len(all_epoch_loss))
plt.figure(1)
plt.plot(x, all_epoch_loss, marker='+', label='Training Loss')  # 绘制训练损失线图
plt.plot(x, all_val_loss, marker='o', color = 'orange', label='Validation Loss')  # 绘制验证损失散点图
plt.title("LOSS")  # 设置标题
plt.xlabel("Epoch")  # 设置X轴标签
plt.ylabel("Loss")  # 设置Y轴标签
plt.legend()  # 显示图例

plt.figure(2)
plt.plot(x, all_epoch_acc, marker='+', label='Training Accuracy')
plt.plot(x, all_val_acc, marker='o', color = 'orange', label='Validation Accuracy')
plt.title("ACCURACY")
plt.xlabel("Epoch")  # 设置X轴标签
plt.ylabel("Accuracy")  # 设置Y轴标签
plt.legend()  # 显示图例

PATH='state_dict_model_multi-class.pth'
#先建立路径
torch.save(model.state_dict(), PATH)

plt.show()

'''optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.L1Loss()

print(val_ds.df)

# 定义训练过程
def train(epoch, model, train_loader, test_loader):
    if epoch == 20:
        print("now!")
    running_loss = 0

    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        with torch.no_grad():
            running_loss += loss.item()

    epoch_loss = running_loss / (len(train_loader.dataset)//batch_size)
    # 测试过程
    test_running_loss = 0
    with torch.no_grad():
        for x, y in test_loader:
            x, y = x.to(device), y.to(device)
            y_pred = model(x)
            loss = loss_fn(y_pred, y)
            test_running_loss += loss.item()
    test_epoch_loss = test_running_loss / (len(test_loader.dataset)//batch_size)
    print('epoch', epoch,
          'train_loss', round(epoch_loss, 3),
          'val_loss', round(test_epoch_loss, 3))
    return epoch_loss, test_epoch_loss

epoch = 20
all_epoch_loss = []
all_val_loss = []
for i in range(epoch):
    epoch_loss, test_epoch_loss = train(i, model, train_loader, val_loader)
    all_epoch_loss.append(epoch_loss)
    all_val_loss.append(test_epoch_loss)

x = range(len(all_epoch_loss))
plt.figure(1)
plt.plot(x, all_epoch_loss, marker='+', label='Training Loss')  # 绘制训练损失线图
plt.plot(x, all_val_loss, marker='o', color = 'orange', label='Validation Loss')  # 绘制验证损失散点图
plt.title("LOSS")  # 设置标题
plt.xlabel("Epoch")  # 设置X轴标签
plt.ylabel("Loss")  # 设置Y轴标签
plt.legend()  # 显示图例

PATH='state_dict_model_regress.pth'
#先建立路径
torch.save(model.state_dict(), PATH)

plt.show()'''